from python_a2a import A2AClient, Message, TextContent, MessageRole, run_server, ErrorContent
import re 
import os 
import requests
from typing import Optional
from python_a2a.server import A2AServer
from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv()) 
class DeepSeekStockAssistant:
    """使用 DeepSeek 模型的股票信息助手。""" 
    
    def __init__(self, api_key: str, duckduckgo_endpoint: str, yfinance_endpoint: str): 
        self.api_key = api_key
        self.api_base = "https://api.deepseek.com/v1" 
        self.duckduckgo_client = A2AClient(duckduckgo_endpoint) 
        self.yfinance_client = A2AClient(yfinance_endpoint) 
        
        # 系统提示词
        self.system_prompt = (
            "你是一个有用的金融助手，可以帮助用户获取股票信息。\n"
            "你从用户查询中提取公司名称以查找股票代码和价格。\n"
            "如果用户询问股票价格，你会协调专门的代理来获取准确信息。"
        )

    def handle_message(self, message: Message) -> Message:
        try:
            if hasattr(message.content, 'text'): 
                text = message.content.text.lower() 
                
                # 股票查询
                stock_keywords = ["stock", "price", "trading", "share", "ticker", "股价", "股票", "价格", "行情"]
                if any(keyword in text for keyword in stock_keywords):
                    return self._get_stock_info_sync(message) 
            
            # 对于所有其他消息，使用 DeepSeek 处理
            return self._call_deepseek_model_sync(message)
        except Exception as e:
            return Message(
                content=TextContent(text=f"抱歉，处理您的消息时遇到了问题：{str(e)}"),
                role=MessageRole.AGENT,
                parent_message_id=getattr(message, 'message_id', None),
                conversation_id=getattr(message, 'conversation_id', None)
            )

    def _call_deepseek_model_sync(self, message: Message) -> Message:
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }
        
        # 构建消息历史
        messages = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": message.content.text}
        ]
        
        payload = {
            "model": "deepseek-chat",
            "messages": messages,
            "temperature": 0.7,
            "max_tokens": 1000
        }
        
        # 发送请求到 DeepSeek API
        response = requests.post(
            f"{self.api_base}/chat/completions",
            headers=headers,
            json=payload,
            timeout=30
        )
        
        if response.status_code != 200:
            return Message(
                content=TextContent(text=f"抱歉，AI 服务暂时不可用。错误代码: {response.status_code}"),
                role=MessageRole.AGENT,
                parent_message_id=getattr(message, 'message_id', None),
                conversation_id=getattr(message, 'conversation_id', None)
            )
        
        result = response.json()
        ai_response = result["choices"][0]["message"]["content"]
        
        return Message(
            content=TextContent(text=ai_response),
            role=MessageRole.AGENT,
            parent_message_id=getattr(message, 'message_id', None),
            conversation_id=getattr(message, 'conversation_id', None)
        )
    
    def _get_stock_info_sync(self, message: Message) -> Message:
        """同步处理股票信息查询。""" 
        # 首先，使用 DeepSeek 提取公司名称
        company_extraction_prompt = (
            f"从以下查询中提取公司名称：'{message.content.text}'。"
            "仅返回公司名称，不返回其他任何内容。"
        )
        
        extraction_message = Message(
            content=TextContent(text=company_extraction_prompt),
            role=MessageRole.USER,
            parent_message_id=getattr(message, 'message_id', None),
            conversation_id=getattr(message, 'conversation_id', None)
        )
        
        openai_response = self._call_deepseek_model_sync(extraction_message)
        
        # 检查响应类型并提取公司名称
        company_name = ""
        if hasattr(openai_response.content, 'text'):
            company_name = openai_response.content.text.strip().strip('"\'.,')
        
        # 如果无法提取公司名称，尝试从原始消息中提取
        if not company_name or len(company_name) < 2:
            company_name = self._extract_company_name(message.content.text)
        
        if not company_name:
            return Message( 
                content=TextContent(text="无法识别公司名称。请明确指定公司，例如：'苹果的股价是多少？'"), 
                role=MessageRole.AGENT, 
                parent_message_id=getattr(message, 'message_id', None), 
                conversation_id=getattr(message, 'conversation_id', None)
            )
        
        print(f"提取的公司名称: {company_name}")
        
        # 步骤 1：从 DuckDuckGo 代理获取股票代码
        ticker_message = Message( 
            content=TextContent(text=f"{company_name}的股票代码是什么？"), 
            role=MessageRole.USER,
            parent_message_id=getattr(message, 'message_id', None),
            conversation_id=getattr(message, 'conversation_id', None)
        ) 
        
        ticker_response = self.duckduckgo_client.send_message(ticker_message)
        
        
        ticker_text = ""
        if hasattr(ticker_response.content, 'text'):
            ticker_text = ticker_response.content.text
        elif hasattr(ticker_response.content, 'error'):
            ticker_text = f"错误: {getattr(ticker_response.content, 'error', '未知错误')}"
        else:
            ticker_text = str(ticker_response.content)
        
        print(f"股票代码响应: {ticker_text}")
        
        # 提取股票代码
        ticker_match = re.search(r'([A-Z]{1,5})', ticker_text)
        if not ticker_match: 
            # 尝试从常见响应格式中提取
            ticker_match = re.search(r'ticker.*?([A-Z]{1,5})', ticker_text, re.IGNORECASE)
            if not ticker_match:
                ticker_match = re.search(r'symbol.*?([A-Z]{1,5})', ticker_text, re.IGNORECASE)
        
        if not ticker_match: 
            return Message( 
                content=TextContent(text=f"找不到 {company_name} 的股票代码。"), 
                role=MessageRole.AGENT, 
                parent_message_id=getattr(message, 'message_id', None), 
                conversation_id=getattr(message, 'conversation_id', None)
            ) 
        
        ticker = ticker_match.group(1) 
        print(f"提取的股票代码: {ticker}")
        
        # 步骤 2：从 YFinance 代理获取股票价格
        price_message = Message( 
            content=TextContent(text=f"{ticker}的当前价格是多少？"), 
            role=MessageRole.USER,
            parent_message_id=getattr(message, 'message_id', None),
            conversation_id=getattr(message, 'conversation_id', None)
        ) 
        
        price_response = self.yfinance_client.send_message(price_message)
        
        
        # 检查价格响应类型
        price_text = ""
        if hasattr(price_response.content, 'text'):
            price_text = price_response.content.text
        elif hasattr(price_response.content, 'error'):
            price_text = f"错误: {getattr(price_response.content, 'error', '未知错误')}"
        else:
            price_text = str(price_response.content)
        
        # 返回组合信息
        return Message( 
            content=TextContent( 
                text=f"{company_name} ({ticker}) 的信息:\n{price_text}"
                ), 
            role=MessageRole.AGENT, 
            parent_message_id=getattr(message, 'message_id', None), 
            conversation_id=getattr(message, 'conversation_id', None)
        ) 
    
    def _extract_company_name(self, text: str) -> str:
        company_mapping = {
            "apple": "Apple",
            "苹果": "Apple", 
            "microsoft": "Microsoft",
            "微软": "Microsoft",
            "google": "Google",
            "谷歌": "Google",
            "amazon": "Amazon",
            "亚马逊": "Amazon",
            "tesla": "Tesla",
            "特斯拉": "Tesla",
            "meta": "Meta",
            "facebook": "Facebook",
            "脸书": "Facebook"
        }
        
        text_lower = text.lower()
        for key, value in company_mapping.items():
            if key in text_lower:
                return value
        
        words = re.findall(r'\b[A-Z][a-z]+\b', text)
        if words:
            return words[0]
        
        return ""


class StockAssistantServer(A2AServer):
    def __init__(self, assistant):
        super().__init__(url="http://localhost:5005")
        self.assistant = assistant
    
    def handle_message(self, message):
        return self.assistant.handle_message(message)

if __name__ == "__main__":
    api_key = os.environ.get("DEEPSEEK_API_KEY")  
    if not api_key:
        print("错误: 需要设置 DEEPSEEK_API_KEY 环境变量")
        exit(1)
    
    assistant = DeepSeekStockAssistant(
        api_key= api_key,
        duckduckgo_endpoint="http://localhost:5003",
        yfinance_endpoint="http://localhost:5004"
    )
    
    run_server(StockAssistantServer(assistant), port=5005)